#include <linux/module.h>
#include <linux/pci.h>
#include <linux/cdev.h>
#include <linux/fs.h>
#include <linux/device.h>
#include <linux/uaccess.h>
#include <linux/io.h>
#include <linux/interrupt.h>
#include <linux/slab.h>
#include <linux/wait.h>
#include <linux/poll.h>
#include <linux/msi.h>
#include <linux/mutex.h>

#define DRV_NAME "pcie_cdev_drv"
#define BAR_NUM 0           // 使用的 BAR 编号
/*
 * 三种中断类型的上限值：
 *  1. MSI-X: PCI 3.0+ 规范支持的最大向量数 (2048)
 *  2. MSI:   PCI 规范支持的最大向量数 (32)
 *  3. Legacy: 传统PCI设备支持的最大中断数 (1)
 * 注意：实际可用数量受设备能力和系统资源限制
 */
#define MAX_IRQS 2048       // 系统支持的最大中断资源数量
#define IRQ_NAME "pcie_cdev_irq"
#define MAX_DEVICES 1       // 支持的最大设备数

// 中断模式枚举
enum irq_mode {
    IRQ_MODE_LEGACY = 0,
    IRQ_MODE_MSI,
    IRQ_MODE_MSIX,
    IRQ_MODE_UNKNOWN
};

// 设备特定数据结构
struct pcie_cdev {
    struct pci_dev *pdev;    // PCI 设备
    void __iomem *regs;      // 映射的寄存器基地址
    resource_size_t reg_len; // 寄存器区域长度
    struct cdev cdev;        // 字符设备
    dev_t devno;             // 设备号
    struct class *cls;       // 设备类
    char device_name[64];    // 唯一设备名称
    
    // 中断相关
    enum irq_mode irq_mode;  // 中断模式
    int num_irqs;            // 实际使用的中断数量
    int irq_vector[MAX_IRQS]; // 中断向量
    atomic_t irq_count;      // 中断计数器
    wait_queue_head_t irq_queue; // 中断等待队列
    struct msix_entry msix_entries[MAX_IRQS]; // MSI-X 表项
    
    // 同步机制
    struct mutex lock;       // 设备锁
};

// 文件操作函数声明
static int pcie_cdev_open(struct inode *inode, struct file *filp);
static int pcie_cdev_release(struct inode *inode, struct file *filp);
static ssize_t pcie_cdev_read(struct file *filp, char __user *buf, 
                             size_t count, loff_t *ppos);
static ssize_t pcie_cdev_write(struct file *filp, const char __user *buf, 
                              size_t count, loff_t *ppos);
static long pcie_cdev_ioctl(struct file *filp, unsigned int cmd, 
                          unsigned long arg);
static unsigned int pcie_cdev_poll(struct file *filp, poll_table *wait);
static irqreturn_t pcie_cdev_isr(int irq, void *dev_id);

// 文件操作结构体
static const struct file_operations pcie_cdev_fops = {
    .owner = THIS_MODULE,
    .open = pcie_cdev_open,
    .release = pcie_cdev_release,
    .read = pcie_cdev_read,
    .write = pcie_cdev_write,
    .unlocked_ioctl = pcie_cdev_ioctl,
    .poll = pcie_cdev_poll,
    .llseek = no_llseek, // 不支持 seek
};

// PCI 设备ID表
static const struct pci_device_id pcie_cdev_ids[] = {
    // { PCI_DEVICE(0x10ee, 0x9021) }, // Xilinx 示例设备
    { PCI_DEVICE(0x0755, 0x0755) },
    { 0, }  // 结束标记
};
MODULE_DEVICE_TABLE(pci, pcie_cdev_ids);

// 打开设备
static int pcie_cdev_open(struct inode *inode, struct file *filp)
{
    struct pcie_cdev *dev;
    
    // 获取与 inode 关联的字符设备
    dev = container_of(inode->i_cdev, struct pcie_cdev, cdev);
    
    // 检查设备是否有效
    if (!dev || !dev->pdev) {
        return -ENODEV;
    }
    
    // 记录设备指针
    filp->private_data = dev;
    
    // 尝试锁定设备
    if (mutex_lock_interruptible(&dev->lock)) {
        return -ERESTARTSYS;
    }
    
    // 简单的打开计数（实际应用中可能需要更复杂的资源管理）
    dev_info(&dev->pdev->dev, "Device opened\n");
    
    mutex_unlock(&dev->lock);
    return 0;
}

// 关闭设备
static int pcie_cdev_release(struct inode *inode, struct file *filp)
{
    struct pcie_cdev *dev = filp->private_data;
    
    if (dev) {
        mutex_lock(&dev->lock);
        dev_info(&dev->pdev->dev, "Device closed\n");
        mutex_unlock(&dev->lock);
    }
    
    return 0;
}

// 读取设备
static ssize_t pcie_cdev_read(struct file *filp, char __user *buf, 
                             size_t count, loff_t *ppos)
{
    struct pcie_cdev *dev = filp->private_data;
    size_t bytes_read = 0;
    ssize_t ret = 0;
    
    if (!dev || !dev->regs || !buf) {
        return -EINVAL;
    }
    
    // 锁保护
    if (mutex_lock_interruptible(&dev->lock)) {
        return -ERESTARTSYS;
    }
    
    // 检查偏移量是否有效
    if (*ppos >= dev->reg_len) {
        ret = 0; // EOF
        goto unlock;
    }
    
    // 计算实际可读取字节数
    bytes_read = min(count, (size_t)(dev->reg_len - *ppos));
    
    // 从寄存器复制到用户空间
    if (copy_to_user(buf, dev->regs + *ppos, bytes_read)) {
        ret = -EFAULT;
        goto unlock;
    }
    
    // 更新文件位置
    *ppos += bytes_read;
    ret = bytes_read;
    
    dev_dbg(&dev->pdev->dev, "Read %zu bytes from offset %lld\n", 
           bytes_read, *ppos - bytes_read);

unlock:
    mutex_unlock(&dev->lock);
    return ret;
}

// 写入设备
static ssize_t pcie_cdev_write(struct file *filp, const char __user *buf, 
                              size_t count, loff_t *ppos)
{
    struct pcie_cdev *dev = filp->private_data;
    size_t bytes_written = 0;
    ssize_t ret = 0;
    
    if (!dev || !dev->regs || !buf) {
        return -EINVAL;
    }
    
    // 锁保护
    if (mutex_lock_interruptible(&dev->lock)) {
        return -ERESTARTSYS;
    }
    
    // 检查偏移量是否有效
    if (*ppos >= dev->reg_len) {
        ret = -EFBIG; // 超出范围
        goto unlock;
    }
    
    // 计算实际可写入字节数
    bytes_written = min(count, (size_t)(dev->reg_len - *ppos));
    
    // 从用户空间复制到寄存器
    if (copy_from_user(dev->regs + *ppos, buf, bytes_written)) {
        ret = -EFAULT;
        goto unlock;
    }
    
    // 更新文件位置
    *ppos += bytes_written;
    ret = bytes_written;
    
    dev_dbg(&dev->pdev->dev, "Wrote %zu bytes to offset %lld\n", 
           bytes_written, *ppos - bytes_written);

unlock:
    mutex_unlock(&dev->lock);
    return ret;
}

// IOCTL 命令定义
#define PC_CHAR_GET_REG_SIZE _IOR('P', 1, size_t)
#define PC_CHAR_SET_REG _IOW('P', 2, struct reg_data)
#define PC_CHAR_GET_REG _IOR('P', 3, struct reg_data)
#define PC_CHAR_ENABLE_IRQ _IO('P', 4)
#define PC_CHAR_DISABLE_IRQ _IO('P', 5)
#define PC_CHAR_GET_IRQ_MODE _IOR('P', 6, int)
#define PC_CHAR_RESET _IO('P', 7)

struct reg_data {
    u32 offset;
    u32 value;
};

// IOCTL 处理
static long pcie_cdev_ioctl(struct file *filp, unsigned int cmd, 
                          unsigned long arg)
{
    struct pcie_cdev *dev = filp->private_data;
    struct reg_data rd;
    long ret = 0;
    
    if (!dev || !dev->regs) {
        return -ENODEV;
    }
    
    // 锁保护
    if (mutex_lock_interruptible(&dev->lock)) {
        return -ERESTARTSYS;
    }
    
    switch (cmd) {
    case PC_CHAR_GET_REG_SIZE:
        ret = put_user(dev->reg_len, (size_t __user *)arg);
        break;
        
    case PC_CHAR_SET_REG:
        if (copy_from_user(&rd, (void __user *)arg, sizeof(rd))) {
            ret = -EFAULT;
            break;
        }
        
        // 验证偏移量有效性
        if (rd.offset >= dev->reg_len - sizeof(u32)) {
            ret = -EINVAL;
            break;
        }
        
        iowrite32(rd.value, dev->regs + rd.offset);
        dev_dbg(&dev->pdev->dev, "Set register 0x%x to 0x%x\n", 
               rd.offset, rd.value);
        break;
        
    case PC_CHAR_GET_REG:
        if (copy_from_user(&rd, (void __user *)arg, sizeof(rd))) {
            ret = -EFAULT;
            break;
        }
        
        // 验证偏移量有效性
        if (rd.offset >= dev->reg_len - sizeof(u32)) {
            ret = -EINVAL;
            break;
        }
        
        rd.value = ioread32(dev->regs + rd.offset);
        if (copy_to_user((void __user *)arg, &rd, sizeof(rd))) {
            ret = -EFAULT;
            break;
        }
        dev_dbg(&dev->pdev->dev, "Read register 0x%x: 0x%x\n", 
               rd.offset, rd.value);
        break;
        
    case PC_CHAR_ENABLE_IRQ:
        // 启用中断 - 实际实现取决于硬件
        iowrite32(0x1, dev->regs + 0x8); // 假设中断控制寄存器在偏移0x8处
        dev_info(&dev->pdev->dev, "Interrupts enabled\n");
        break;
        
    case PC_CHAR_DISABLE_IRQ:
        // 禁用中断
        iowrite32(0x0, dev->regs + 0x8);
        dev_info(&dev->pdev->dev, "Interrupts disabled\n");
        break;
        
    case PC_CHAR_GET_IRQ_MODE:
        ret = put_user((int)dev->irq_mode, (int __user *)arg);
        break;
        
    case PC_CHAR_RESET:
        // 设备复位 - 示例实现
        // iowrite32(0x1, dev->regs + 0x0); // 写控制寄存器复位位
        // msleep(10); // 等待复位完成
        // iowrite32(0x0, dev->regs + 0x0);
        // dev_info(&dev->pdev->dev, "Device reset\n");
        break;
        
    default:
        ret = -ENOTTY;
        break;
    }
    
    mutex_unlock(&dev->lock);
    return ret;
}

// Poll 函数 - 用于等待中断
static unsigned int pcie_cdev_poll(struct file *filp, poll_table *wait)
{
    struct pcie_cdev *dev = filp->private_data;
    unsigned int mask = 0;
    
    if (!dev) {
        return POLLERR;
    }
    
    // 添加到等待队列
    poll_wait(filp, &dev->irq_queue, wait);
    
    // 检查中断状态
    if (atomic_read(&dev->irq_count) > 0) {
        mask |= POLLIN | POLLRDNORM; // 数据可读
    }
    
    return mask;
}

// 中断服务例程
static irqreturn_t pcie_cdev_isr(int irq, void *dev_id)
{
    struct pcie_cdev *dev = dev_id;
    int irq_idx = -1;
    int i;
    u32 status = 0;
    
    if (!dev || !dev->regs) {
        return IRQ_NONE;
    }
    
    // 确定中断索引 (仅对MSI-X/MSI重要)
    if (dev->irq_mode == IRQ_MODE_MSIX || dev->irq_mode == IRQ_MODE_MSI) {
        for (i = 0; i < dev->num_irqs; i++) {
            if (dev->irq_vector[i] == irq) {
                irq_idx = i;
                break;
            }
        }
    } else {
        irq_idx = 0; // Legacy中断只有一个
    }
    dev_info(&dev->pdev->dev, "Interrupt %d handled (vector %d)\n", 
             irq, irq_idx);
    if (irq_idx < 0) {
        pr_err("Unknown interrupt %d\n", irq);
        return IRQ_NONE;
    }
    
    // 1. 检查中断状态
    status = ioread32(dev->regs + 0x4); // 状态寄存器
    if (!(status & (1 << irq_idx))) {
        return IRQ_NONE; // 不是我们的中断
    }
    
    // 2. 清除中断标志
    iowrite32(status | (1 << irq_idx), dev->regs + 0x4);
    
    // 3. 更新中断计数器
    atomic_inc(&dev->irq_count);
    
    // 4. 唤醒等待进程
    wake_up_interruptible(&dev->irq_queue);
    
    // 5. 记录中断事件
    dev_dbg(&dev->pdev->dev, "Interrupt %d handled (vector %d)\n", 
           irq, irq_idx);
    
    return IRQ_HANDLED;
}

// 初始化中断系统
static int init_interrupts(struct pcie_cdev *dev)
{
    struct pci_dev *pdev = dev->pdev;
    int ret = 0, i;
    int max_irqs = MAX_IRQS;

    // 尝试使用 MSI-X
    dev->irq_mode = IRQ_MODE_UNKNOWN;
    ret = pci_alloc_irq_vectors(pdev, 1, max_irqs, PCI_IRQ_MSIX);
    if (ret > 0) {
        dev->num_irqs = ret;
        dev->irq_mode = IRQ_MODE_MSIX;
        dev_info(&pdev->dev, "Using MSI-X with %d vectors\n", dev->num_irqs);
        goto setup_irqs;
    }

    // 尝试使用 MSI
    ret = pci_alloc_irq_vectors(pdev, 1, max_irqs, PCI_IRQ_MSI);
    if (ret > 0) {
        dev->num_irqs = ret;
        dev->irq_mode = IRQ_MODE_MSI;
        dev_info(&pdev->dev, "Using MSI with %d vectors\n", dev->num_irqs);
        goto setup_irqs;
    }

    // 使用 Legacy 中断
    ret = pci_alloc_irq_vectors(pdev, 1, 1, PCI_IRQ_LEGACY);
    if (ret < 0) {
        dev_err(&pdev->dev, "Failed to allocate IRQ vectors (%d)\n", ret);
        return ret;
    }

    dev->num_irqs = 1;
    dev->irq_mode = IRQ_MODE_LEGACY;
    dev_info(&pdev->dev, "Using legacy INTx interrupt\n");

setup_irqs:
    // 获取中断号并注册处理程序
    for (i = 0; i < dev->num_irqs; i++) {
        dev->irq_vector[i] = pci_irq_vector(pdev, i);
        if (dev->irq_vector[i] < 0) {
            dev_err(&pdev->dev, "Invalid IRQ vector for index %d\n", i);
            ret = -EINVAL;
            goto free_irq_vectors;
        }

        // 注册中断处理程序
        ret = request_irq(dev->irq_vector[i], pcie_cdev_isr, 
                          (dev->irq_mode == IRQ_MODE_LEGACY) ? IRQF_SHARED : 0,
                          IRQ_NAME, dev);
        if (ret) {
            dev_err(&pdev->dev, "Failed to request IRQ %d for vector %d\n",
                    dev->irq_vector[i], i);
            goto free_irqs;
        }
    }

    return 0;

free_irqs:
    // 清理已分配的中断
    while (i-- > 0) {
        free_irq(dev->irq_vector[i], dev);
    }

free_irq_vectors:
    pci_free_irq_vectors(pdev);
    return ret;
}

// 清理中断系统
static void cleanup_interrupts(struct pcie_cdev *dev)
{
    struct pci_dev *pdev = dev->pdev;
    int i;

    for (i = 0; i < dev->num_irqs; i++) {
        free_irq(dev->irq_vector[i], dev);
    }

    pci_free_irq_vectors(pdev);

    dev->num_irqs = 0;
    dev->irq_mode = IRQ_MODE_UNKNOWN;
}

// PCI 设备探测函数
static int pcie_cdev_probe(struct pci_dev *pdev, const struct pci_device_id *id)
{
    struct pcie_cdev *dev;
    int ret;
    dev_t devno;

    dev_info(&pdev->dev, "Device: %04x:%04x at %s\n", 
             pdev->vendor, pdev->device, pci_name(pdev));

    // 0. 检查设备是否已初始化
    if (pci_get_drvdata(pdev)) {
        dev_info(&pdev->dev, "Device already initialized\n");
        return -EBUSY;
    }

    // 1. 分配设备结构
    dev = devm_kzalloc(&pdev->dev, sizeof(*dev), GFP_KERNEL);
    if (!dev)
        return -ENOMEM;
    
    // 2. 保存PCI设备
    dev->pdev = pdev;
    pci_set_drvdata(pdev, dev);
    
    // 3. 生成唯一设备名称（使用PCI位置）
    snprintf(dev->device_name, sizeof(dev->device_name), "pcie_cdev_%04x:%02x:%02x.%d",
             pci_domain_nr(pdev->bus), pdev->bus->number,
             PCI_SLOT(pdev->devfn), PCI_FUNC(pdev->devfn));
    
    // 4. 初始化互斥锁
    mutex_init(&dev->lock);
    
    // 5. 初始化等待队列
    init_waitqueue_head(&dev->irq_queue);
    atomic_set(&dev->irq_count, 0);
    
    // 6. 启用PCI设备
    ret = pci_enable_device(pdev);
    if (ret) {
        dev_err(&pdev->dev, "Failed to enable PCI device\n");
        goto err_free_dev;
    }
    
    // 7. 请求内存区域
    ret = pci_request_regions(pdev, dev->device_name);
    if (ret) {
        dev_err(&pdev->dev, "Failed to request regions\n");
        goto err_disable_dev;
    }
    
    // 8. 检查资源长度
    dev->reg_len = pci_resource_len(pdev, BAR_NUM);
    if (dev->reg_len == 0) {
        dev_err(&pdev->dev, "Invalid BAR size\n");
        ret = -EIO;
        goto err_release_regions;
    }
    
    // 9. 映射寄存器区域
    dev->regs = pci_ioremap_bar(pdev, BAR_NUM);
    if (!dev->regs) {
        dev_err(&pdev->dev, "Failed to map registers\n");
        ret = -ENOMEM;
        goto err_release_regions;
    }
    
    // 10. 设置DMA掩码
    if (pci_set_dma_mask(pdev, DMA_BIT_MASK(64))) {
        if (pci_set_dma_mask(pdev, DMA_BIT_MASK(32))) {
            dev_err(&pdev->dev, "No suitable DMA available\n");
            ret = -EIO;
            goto err_unmap_regs;
        }
    }
    pci_set_master(pdev);
    
    // 11. 初始化中断系统
    ret = init_interrupts(dev);
    if (ret) {
        dev_err(&pdev->dev, "Failed to initialize interrupts\n");
        goto err_unmap_regs;
    }
    
    // 12. 为每个设备分配设备号
    ret = alloc_chrdev_region(&devno, 0, MAX_DEVICES, dev->device_name);
    if (ret) {
        dev_err(&pdev->dev, "Failed to allocate devno\n");
        goto err_cleanup_irq;
    }
    dev->devno = devno;
    
    // 13. 创建设备类
    dev->cls = class_create(THIS_MODULE, "pcie_cdev_class");
    if (IS_ERR(dev->cls)) {
        ret = PTR_ERR(dev->cls);
        dev_err(&pdev->dev, "Failed to create class\n");
        goto err_unregister_dev;
    }
    
    // 14. 初始化字符设备
    cdev_init(&dev->cdev, &pcie_cdev_fops);
    dev->cdev.owner = THIS_MODULE;
    
    // 15. 添加字符设备到系统
    ret = cdev_add(&dev->cdev, dev->devno, 1);
    if (ret) {
        dev_err(&pdev->dev, "Failed to add cdev\n");
        goto err_class_destroy;
    }
    
    // 16. 创建设备节点（使用唯一名称）
    device_create(dev->cls, NULL, dev->devno, NULL, "%s", dev->device_name);
    
    dev_info(&pdev->dev, "PCIe character device %s initialized\n", dev->device_name);
    dev_info(&pdev->dev, "BAR%d size: %pa bytes\n", BAR_NUM, &dev->reg_len);
    dev_info(&pdev->dev, "Using %s with %d interrupt vectors\n",
             dev->irq_mode == IRQ_MODE_MSIX ? "MSI-X" :
             dev->irq_mode == IRQ_MODE_MSI ? "MSI" : "Legacy",
             dev->num_irqs);
             
    return 0;

// 错误处理路径
err_class_destroy:
    class_destroy(dev->cls);
err_unregister_dev:
    unregister_chrdev_region(dev->devno, MAX_DEVICES);
err_cleanup_irq:
    cleanup_interrupts(dev);
err_unmap_regs:
    if (dev->regs) {
        pci_iounmap(pdev, dev->regs);
    }
err_release_regions:
    pci_release_regions(pdev);
err_disable_dev:
    pci_disable_device(pdev);
err_free_dev:
    // devm 会自动释放内存
    return ret;
}

// PCI 设备移除函数
static void pcie_cdev_remove(struct pci_dev *pdev)
{
    struct pcie_cdev *dev = pci_get_drvdata(pdev);
    
    if (!dev) {
        return;
    }
    
    // 1. 销毁设备节点
    device_destroy(dev->cls, dev->devno);
    
    // 2. 删除字符设备
    cdev_del(&dev->cdev);
    
    // 3. 销毁设备类
    class_destroy(dev->cls);
    
    // 4. 释放设备号
    unregister_chrdev_region(dev->devno, 1);
    
    // 5. 清理中断系统
    cleanup_interrupts(dev);
    
    // 6. 取消寄存器映射
    if (dev->regs) {
        pci_iounmap(pdev, dev->regs);
    }
    
    // 7. 释放PCI区域
    pci_release_regions(pdev);
    
    // 8. 禁用PCI设备
    pci_disable_device(pdev);
    
    // 9. 销毁互斥锁
    mutex_destroy(&dev->lock);
    
    dev_info(&pdev->dev, "PCIe character device %s removed\n", dev->device_name);
}

// PCI 驱动结构
static struct pci_driver pcie_cdev_driver = {
    .name = DRV_NAME,
    .id_table = pcie_cdev_ids,
    .probe = pcie_cdev_probe,
    .remove = pcie_cdev_remove,
};

// 模块初始化和退出
static int __init pcie_cdev_init(void)
{
    int ret;
    
    ret = pci_register_driver(&pcie_cdev_driver);
    if (ret) {
        pr_err("Failed to register PCI driver\n");
    } else {
        pr_info("PCIe character device driver loaded\n");
    }
    
    return ret;
}

static void __exit pcie_cdev_exit(void)
{
    pci_unregister_driver(&pcie_cdev_driver);
    pr_info("PCIe character device driver unloaded\n");
}

module_init(pcie_cdev_init);
module_exit(pcie_cdev_exit);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Jason");
MODULE_DESCRIPTION("PCIe Character Device Driver");
MODULE_VERSION("0.1");
